from z_common import *
from stock_list import *
from finance_list import *
from z_tdx_favlist import *
import pandas as pd
import numpy as np
import os

filename = './analyse_data/沪深Ａ股20210710.txt'

if __name__ == '__main__' : 
    dfi = getStockInfoDataFrame(filename, ['代码', '名称', '现价', '细分行业'])

    # for test
    # dfi = dfi.iloc[0:10]

    dfi['扣非净利润(万元)'] = None;
    dfi['扣非净利润增长率(%)'] = None;
    dfi['扣非净利润(季)(万元)'] = None;
    dfi['扣非净利润(季)增长率(%)'] = None;
    dfi['EPS'] = None;
    dfi['PE(静态)'] = None;
    dfi['PEG'] = None;
    dfi['ROE(%)'] = None;
    dfi_len = len(dfi)
    for i in range(dfi_len):
        if i % int(dfi_len / 10) == 0:
            print('...' + str((int)(i / int(dfi_len / 10) * 10)) + '%...', end='')

        code = dfi['代码'].iloc[i]
        findf = financeDataFrame(code)

        profit0 = finone(findf, '净利润(扣除非经常性损益后)(万元)', '12-31')
        profit1 = finone(findf, '净利润(扣除非经常性损益后)(万元)', '12-31', 1)
        dfi['扣非净利润(万元)'].iloc[i] = profit0
        if not pd.isna(profit0) and not pd.isna(profit1):
            dfi['扣非净利润增长率(%)'].iloc[i] = round(( profit0/profit1  - 1) * 100, 2)
        
        profitS0 = finone(findf, '净利润(扣除非经常性损益后)(万元)', 'fixed-last-season')
        profitS1 = finone(findf, '净利润(扣除非经常性损益后)(万元)', 'fixed-last-season', 1)
        dfi['扣非净利润(季)(万元)'].iloc[i] = profitS0
        if not pd.isna(profitS0) and not pd.isna(profitS1):
            dfi['扣非净利润(季)增长率(%)'].iloc[i] = round(( profitS0/profitS1  - 1) * 100, 2)
        
        dfi['EPS'].iloc[i] = finone(findf, '基本每股收益(元)', '12-31')

        dfi['ROE(%)'].iloc[i] = finone(findf, '净资产收益率加权(%)', '12-31')
    
    dfi['PE(静态)'] = dfi['现价'] / dfi['EPS']
    dfi['PE(静态)'].apply(np.round, 2)
    dfi['PEG'] = dfi['PE(静态)'] / dfi['扣非净利润增长率(%)']
    dfi['PEG'].apply(np.round, 2)

    # 全部(根据扣非净利润排序)
    dfi.sort_values('扣非净利润(万元)', ascending=False, inplace=True)
    reportname = filename.replace("/analyse_data/", "/report/").replace(".txt", ".csv")
    dfi.to_csv(reportname, index=False, encoding='ANSI')

    # 筛选
    selname = reportname.replace(".csv", ".sel.csv")
    dfi = dfi[
          (dfi["扣非净利润(季)(万元)"] > 0) \
        & (dfi["ROE(%)"] > 10.0) \
        & (dfi["扣非净利润(季)增长率(%)"] > 15.0) \
        & (dfi["扣非净利润(季)增长率(%)"] >= dfi["ROE(%)"]) \
        & (dfi["PEG"] > 0.0) \
        & (dfi["PEG"] < 1.0)]
    dfi.to_csv(selname, index=False, encoding='ANSI')

    # 输出通达信自选股列表
    ebkname = reportname.replace(".csv", ".sel.ebk")
    output_ebk(dfi, ebkname)